import os
import sys
import numpy as np
import glob
sys.path.append(os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.realpath(__file__))))), 'util', 'file'))
from nii import NiiFile

class VerSeDataset:
    """
    CT images and their masks, see https://gitee.com/yawei-gitee/medical-datasets-ongoing/tree/master/vertebrae. 
    
    * Samples: 374
    * Original label: masks
    * Task: segmentation
    """
    def __init__(self, givenDirOfDataset: str, givenLoadType: str = 'data'):
        self.__name = 'VerSe'
        self.__dir = givenDirOfDataset
        self.__numberOfSamples = 374
        self.__samples = None
        self.__labels = None
        self.__pathListOfSamples = self.__queryPathListOfSamples()
        self.__pathListOfLabels = self.__queryPathListOfLabels()
        self.__randomSeed = 0
        self.__loadType = givenLoadType
        self.__init()
        pass
    
    def __init(self):
        if self.__loadType == 'data':
            self.__initData()
            self.__shuffle()
        pass
    
    def __shuffle(self):
        self.__shuffleSamples()
        self.__shuffleLabels()
        pass
    
    def __shuffleSamples(self):
        """
        re-organize the samples randomly
        """
        np.random.seed(self.__randomSeed)
        self.__samples = np.random.permutation(self.__samples)
        pass
    
    def __shuffleLabels(self):
        """
        re-organize the samples randomly
        """
        np.random.seed(self.__randomSeed)
        self.__labels = np.random.permutation(self.__labels)
        pass
    
    def __queryPathListOfSamples(self):
        imgPaths = glob.glob(f"{self.__dir}{os.sep}rawdata{os.sep}*{os.sep}*.nii.gz")
        return imgPaths
    
    def __queryPathListOfLabels(self):
        imgPaths = glob.glob(f"{self.__dir}{os.sep}derivatives{os.sep}*{os.sep}*.nii.gz")
        return imgPaths

    def __initSamples(self):
        imgList = []
        for imgPath in self.__pathListOfSamples:
            imgData = NiiFile(givenFilePath=imgPath).data
            imgList.append(imgData)
        self.__samples = np.array(imgList)
        pass

    def __initLabels(self):
        imgList = []
        for imgPath in self.__pathListOfLabels:
            imgData = NiiFile(givenFilePath=imgPath).data
            imgList.append(imgData)
        self.__labels = np.array(imgList)
        pass

    def __initData(self):
        self.__initSamples()
        self.__initLabels()
        pass

    def querySample(self, givenSampleId):
        return self.__samples[givenSampleId]
    
    def queryLabel(self, givenSampleId):
        return self.__labels[givenSampleId]
    
    def querySampleList(self, givenSampleIds: list):
        assert(isinstance(givenSampleIds, list))
        sampleList = [self.__samples[id] for id in givenSampleIds] 
        return np.array(sampleList)
    
    def getSamples(self):
        return self.__samples

    def getLabels(self):
        return self.__labels
    
    def showSample(self, givenId: int):
        imgPath = self.__pathListOfSamples[givenId]
        NiiFile(givenFilePath=imgPath).show()
        pass

    def showLabel(self, givenId: int):
        imgPath = self.__pathListOfLabels[givenId]
        NiiFile(givenFilePath=imgPath).show()
        pass

    @property
    def numberOfSamples(self):
        return self.__numberOfSamples
    
    @property
    def filePathOfSameples(self):
        return self.__pathListOfSamples
    
    @property
    def filePathOfLabels(self):
        return self.__pathListOfLabels

    @property
    def name(self):
        return self.__name
    pass










class Test:
    def execute(self):
        vsd = VerSeDataset(givenDirOfDataset='E:\\images\\vertebrae\\verse\\basic', givenLoadType='path') # 'C:\\Users\\TFBOYS\\Documents\\data\\verse'
        vsd.showLabel(givenId=10)
        vsd.showSample(givenId=10)
        pass
    pass


if __name__ == '__main__':
    Test().execute()





